package com.qingzhuge.limiter.support;

import com.google.gson.Gson;
import com.qingzhuge.limiter.dto.RateLimiterConfig;
import com.qingzhuge.limiter.dto.RateLimiterForm;
import com.qingzhuge.limiter.exception.RateLimiterException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.util.StringUtils;

import java.util.*;

/**
 * @author zeroxiao
 * @date 2019/9/5 18:39
 */
@Slf4j
public class RateLimiterHandler implements InitializingBean {
    @Autowired
    private RedisTemplate<String, Long> redisTemplate;
    @Autowired
    private StringRedisTemplate stringRedisTemplate;
    /**
     * 缓存回调，为执行限流时调用业务方法
     */
    private Map<Class<? extends RateLimiterCallback>, RateLimiterCallback> callbackMap = new HashMap<>(16);
    /**
     * 业务回调实现集合，当有多个业务实现限流时，会全布注入集合
     */
    @Autowired(required = false)
    private List<RateLimiterCallback> callbackList = new ArrayList<>();
    /**
     * Lua 脚本
     */
    private DefaultRedisScript<Long> script;
    /**
     * 限流配置 redis key
     */
    private final String KEY_LIMITER_CONFIG = "limiter:config:%s:%s";

    /**
     * 初始化 lua 脚本
     */
    private void init() {
        script = new DefaultRedisScript<>();
        script.setResultType(Long.class);
        script.setLocation(new ClassPathResource("rate_limiter.lua"));
        log.info("RateLimiterHandler[分布式限流处理器]脚本加载完成");
    }

    /**
     * 保存业务限流配置
     */
    @Override
    public void afterPropertiesSet() {
        //初始化 lua 脚本
        this.init();
        //初始化业务限流配置
        if (callbackList.size() > 0) {
            log.debug("业务限流配置初始化.....");
            callbackList.forEach((callback) -> {
                Class<? extends RateLimiterCallback> aClass = callback.getClass();
                //把接口放入map中，方便回调
                this.callbackMap.put(aClass, callback);
                List<RateLimiterConfig> initList = callback.init();
                if (initList != null) {
                    initList.forEach(this::saveConfig);
                }
            });
        }
    }

    private void saveConfig(RateLimiterConfig config) {
        // 业务限流key
        String key = config.getKey();
        //业务模块，业务分类
        String module = config.getModule();
        stringRedisTemplate.opsForValue().set(String.format(KEY_LIMITER_CONFIG, module, key), new Gson().toJson(config));
    }

    /**
     * 执行限流业务
     *
     * @param callbackClass 业务回调实现
     * @param form          限流业务
     * @throws RateLimiterException 发生限流时抛出此异常
     */
    public void execute(Class<? extends RateLimiterCallback> callbackClass, RateLimiterForm form) throws RateLimiterException {
        // 限流阈值
        long limit = form.getLimit();
        //小于1时，不限流
        if(limit < 1){
            return;
        }
        //业务模块，业务分类
        String module = form.getModule();
        String key = form.getKey();
        String level = form.getLevel();
        // 限流超时时间,默认1秒
        long expireTime = form.getExpire();
        if (expireTime < 1){
            expireTime =1;
        }
        String configKey;
        String singleton = "singleton";
        if (null == level) {
            level = singleton;
            configKey = String.format("%s:%s",level,key);
        }else{
            configKey = level;
        }
        //从redis中获取配置
        String configJson = stringRedisTemplate.opsForValue().get(String.format(KEY_LIMITER_CONFIG, module, configKey));
        RateLimiterCallback callback;
        RateLimiterConfig limiterConfig = null;
        if (StringUtils.isEmpty(configJson)) {
            if (singleton.equals(level)){
                limiterConfig = RateLimiterConfig
                        .builder()
                        .expire(expireTime)
                        .limit(limit)
                        .module(module)
                        .key(configKey)
                        .build();
            }else {
                callback = callbackMap.get(callbackClass);
                if (null != callback) {
                    limiterConfig = callback.config(form);
                }
            }
            if (null != limiterConfig) {
                saveConfig(limiterConfig);
            }
        } else {
            //解析业务限流配置
            limiterConfig = new Gson().fromJson(configJson, RateLimiterConfig.class);
        }

        if (limiterConfig != null) {
            // 限流阈值
            limit = limiterConfig.getLimit();
            // 限流超时时间
            expireTime = limiterConfig.getExpire();
        }
        //限流业务 redis key
        String limiterKey = String.format("limiter:operation:%s:%s:%s", module, level, key);
        //调用脚本并执行Lua脚本
        List<String> keys = Collections.singletonList(limiterKey);
        Long result = redisTemplate.execute(script, keys, expireTime, limit);
        if (null != result && 0 == result) {
            callback = callbackMap.get(callbackClass);
            String message = String.format("limiter operation,key:%s ,rate limiter exception", key);
            if (callback != null) {
                message = callback.message(key);
            }
            log.debug("触发限流时间[{}]s内允许的请求数[{}],message:[{}]", expireTime, limit, message);
            throw new RateLimiterException(message);
        }
    }
}
